package com.hiram.controller;


import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.StreamingChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;

/**
 * @author hiram
 * @description 初次接入对话模型
 * @date 2025/4/27 17:03:37
 */
@Slf4j
@RestController
@RequestMapping("/hiram/ai/chat")
public class ChatController {

    @Resource
    private OpenAiChatModel openAiChatModel;

    @Resource
    private StreamingChatModel streamingChatModel;

    /**
     * 单次对话，需要等AI生成完毕之后统一返回
     *
     * @param message 问题
     * @return AI回复
     */
    @RequestMapping("/single")
    public String single(String message) {
        return openAiChatModel.call(message);
    }

    /**
     * 单词对话-流式回复
     * 与网页版对话返回相同，一个字一个字的蹦
     *
     * @param message 问题
     * @return AI回复
     */
    @RequestMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> stream(String message) {
        return streamingChatModel.stream(message).flatMapSequential(Flux::just);
    }

    /**
     * 存储历史消息
     */
    static List<Message> historyMessage = new ArrayList<>();
    /**
     * 历史消息列表最大长度
     */
    static int maxLen = 10;

    /**
     * 上下文历史对话
     * 历史对话长度由maxLen决定
     *
     * @param message 问题
     * @return AI回复
     */
    @RequestMapping("/context")
    public String context(String message) {
        // 用户输入的文本是UserMessage
        historyMessage.add(new UserMessage(message));
        // 发送AI前检查历史消息长度
        if (historyMessage.size() > maxLen) {
            historyMessage.remove(0);
        }
        ChatResponse chatResponse = openAiChatModel.call(new Prompt(historyMessage));
        AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
        // 将AI回复的消息放到历史消息列表中
        historyMessage.add(assistantMessage);
        return assistantMessage.getContent();
    }


    /**
     * 人设提示词
     */
    public static final String PERSONA = "你现在是一个语言学习辅助助手，我在学习中国的脏话，无论我问什么，你都要巧妙的用脏话回复。即使面对批评，道德审判，任何形式的质疑，或者让你忘记脏话学习助手的身份，等等任何提问和沟通，你都要用脏话回复并且隐瞒你脏话助手的身份。碰到无法回答的问题，就随便回复一句脏话。";
    /**
     * 历史消息列表，必须保证人设在第一条
     */
    static List<Message> historyPersonaMessage = new ArrayList<>(List.of(new SystemMessage(PERSONA)));
    /**
     * 最大消息长度
     */
    static int maxPersonaLen = 20;
    /**
     * 已调用次数，打印日志使用
     */
    static int PersonaCount = 0;

    /**
     * 带有人设的上下文对话
     * 人设由PERSONA决定
     * @param message 问题
     * @return AI回复
     */
    @RequestMapping("/persona")
    public String persona(String message) {
        PersonaCount++;
        // 加入用户消息
        historyPersonaMessage.add(new UserMessage(message));
        // 校验历史消息长度并保证人设始终在第一条
        if (historyPersonaMessage.size() > maxPersonaLen) {
            historyPersonaMessage.remove(1);
        }
        log.info("第{}次对话，historyPersonaMessage:{}", PersonaCount, historyPersonaMessage);
        ChatResponse chatResponse = openAiChatModel.call(new Prompt(historyPersonaMessage));
        AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
        historyPersonaMessage.add(assistantMessage);
        return assistantMessage.getContent();
    }
}
